{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "737c4e5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import os\n",
    "import requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d1e2cf93",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('sales_textbook.txt'):\n",
    "    url = 'https://huggingface.co/datasets/goendalf666/sales-textbook_for_convincing_and_selling/resolve/main/sales_textbook.txt'\n",
    "    response = requests.get(url)\n",
    "    with open('sales_textbook.txt', 'w') as f:\n",
    "        f.write(response.text)\n",
    "\n",
    "with open('sales_textbook.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()\n",
    "\n",
    "# text[:500]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "98a066e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "batch_size = 4      # x个批次并行训练\n",
    "context_size = 16   # 输入序列长度\n",
    "d_model = 64\n",
    "num_heads = 4       # 多头注意力的头数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tiktoken\n",
    "enc = tiktoken.get_encoding(\"o200k_base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "816cfa1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "76923\n",
      "199853\n",
      "tensor([ 45990,    220,     16,     25,  22478, 158287,    326,  29742,   4788,\n",
      "         81932])\n"
     ]
    }
   ],
   "source": [
    "tokenized_text = enc.encode(text)\n",
    "tokenized_text = torch.tensor(tokenized_text, dtype=torch.long)     # len() 表示用到的文字的token的个数\n",
    "max_token_value = int(torch.max(tokenized_text))                    # 词表里的token的最大值, 词表里的词不一定都用到了\n",
    "\n",
    "# split to train and validation sets\n",
    "train_idx= int(len(tokenized_text) * 0.9)\n",
    "train_data = tokenized_text[:train_idx]\n",
    "val_data = tokenized_text[train_idx:]\n",
    "\n",
    "print(len(tokenized_text))\n",
    "print(max_token_value)\n",
    "print(tokenized_text[:10])  # Print the first 10 tokens to verify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "518cfd7d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([24302, 68156, 63139,  2249])\n",
      "torch.Size([4, 16]) torch.Size([4, 16])\n"
     ]
    }
   ],
   "source": [
    "data = train_data\n",
    "batch_idxs = torch.randint(low = 0, high = len(data) - context_size, size = (batch_size,))\n",
    "\n",
    "x_batch = torch.stack([data[i:i+context_size] for i in batch_idxs])     # tokenized text\n",
    "y_batch = torch.stack([data[i+1:i+1+context_size] for i in batch_idxs]) # y相比x向前一个位置，作为预测目标\n",
    "\n",
    "print(batch_idxs)\n",
    "print(x_batch.shape, y_batch.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8d7c1936",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' challenges and demonstrating your understanding, you establish yourself as a reliable and trustworthy salesperson.\\n'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "df = pd.DataFrame(x_batch.numpy())\n",
    "enc.decode(df.iloc[0].tolist())  # decode the first row of x_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12873f38",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedding(199854, 64)\n",
      "Parameter containing:\n",
      "tensor([[ 0.9746,  0.0257,  1.2428,  ...,  0.8464, -1.5246,  0.5678],\n",
      "        [-0.6309, -0.7373,  1.3936,  ..., -0.4268, -1.2555,  0.4132],\n",
      "        [-1.8322, -0.3063, -0.8633,  ...,  0.3685,  1.7548, -1.1050],\n",
      "        ...,\n",
      "        [ 0.4536, -0.3777,  1.2460,  ...,  0.2020,  0.8204,  1.1518],\n",
      "        [-0.4620, -2.5583, -1.1699,  ...,  0.2427, -0.0967, -0.9821],\n",
      "        [ 1.6256,  1.3452, -1.8877,  ...,  0.3258,  1.2051,  0.5673]],\n",
      "       requires_grad=True)\n",
      "torch.Size([4, 16, 64]) torch.Size([4, 16, 64])\n"
     ]
    }
   ],
   "source": [
    "# Define input embedding table\n",
    "input_embedding_lookup_table = nn.Embedding(num_embeddings=max_token_value + 1, embedding_dim=d_model)\n",
    "\n",
    "x_batch_embedding = input_embedding_lookup_table(x_batch)  # (batch_size, context_size, d_model)\n",
    "y_batch_embedding = input_embedding_lookup_table(y_batch)\n",
    "\n",
    "print(input_embedding_lookup_table)\n",
    "print(input_embedding_lookup_table.weight)\n",
    "print(x_batch_embedding.shape, y_batch_embedding.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1e3375b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        ...,\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.]])\n",
      "torch.Size([16, 64])\n",
      "torch.Size([4, 16, 64])\n"
     ]
    }
   ],
   "source": [
    "# Define posional encoding lookup table\n",
    "import math\n",
    "position_encoding_lookup_table = torch.zeros(context_size, d_model)         # (context_size, d_model)\n",
    "position = torch.arange(0, context_size, dtype=torch.float).unsqueeze(1)    #(context_size,) -> (context_size, 1)\n",
    "\n",
    "print(position_encoding_lookup_table)\n",
    "print(position_encoding_lookup_table.shape)\n",
    "\n",
    "div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))               # (d_model/2,)\n",
    "position_encoding_lookup_table[:, 0::2] = torch.sin(position * div_term)                                 # (context_size, d_model/2)\n",
    "position_encoding_lookup_table[:, 1::2] = torch.cos(position * div_term)                                 # (context_size, d_model/2)\n",
    "position_encoding_lookup_table = position_encoding_lookup_table.unsqueeze(0).expand(batch_size, -1, -1)  # (1, context_size, d_model)\n",
    "\n",
    "# print(position_encoding_lookup_table)\n",
    "print(position_encoding_lookup_table.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "424d2394",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 16, 64]) torch.Size([4, 16, 64])\n"
     ]
    }
   ],
   "source": [
    "x = x_batch_embedding + position_encoding_lookup_table\n",
    "y = y_batch_embedding + position_encoding_lookup_table\n",
    "\n",
    "print(x.shape, y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5724d027",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 16, 64]) torch.Size([4, 16, 64]) torch.Size([4, 16, 64])\n",
      "torch.Size([4, 4, 16, 16]) torch.Size([4, 4, 16, 16]) torch.Size([4, 4, 16, 16])\n"
     ]
    }
   ],
   "source": [
    "# Get Q, K, V\n",
    "Wq = nn.Linear(d_model, d_model)\n",
    "Wk = nn.Linear(d_model, d_model)\n",
    "Wv = nn.Linear(d_model, d_model)\n",
    "\n",
    "Q = Wq(x)   # batch_size 自动计算了\n",
    "K = Wk(x)\n",
    "V = Wv(x)\n",
    "\n",
    "print(Q.shape, K.shape, V.shape)\n",
    "\n",
    "# Create muti heads\n",
    "Q = Q.view(batch_size, context_size, num_heads, d_model//num_heads).permute(0, 2, 1, 3)\n",
    "K = K.view(batch_size, context_size, num_heads, d_model//num_heads).permute(0, 2, 1, 3)\n",
    "V = V.view(batch_size, context_size, num_heads, d_model//num_heads).permute(0, 2, 1, 3)\n",
    "\n",
    "print(Q.shape, K.shape, V.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "57767ba4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 4, 16, 16])\n",
      "torch.Size([4, 16, 64])\n"
     ]
    }
   ],
   "source": [
    "# Apply attention equation\n",
    "output = Q @ K.transpose(-2,-1) / math.sqrt(d_model//num_heads)\n",
    "\n",
    "# mask\n",
    "mask = torch.triu(torch.ones(context_size, context_size), diagonal=1).bool()\n",
    "output = output.masked_fill(mask, float('-inf'))\n",
    "# pd.DataFrame(output[0,0].detach().numpy())\n",
    "\n",
    "# softmax\n",
    "output = F.softmax(output, dim=-1)      # 最后一个维度是 word 的 token embedding, 永远是对这个维度操作\n",
    "output = output @ V                     # (batch_size, num_heads, context_size, d_model) \n",
    "\n",
    "print(output.shape)\n",
    "\n",
    "# Concatenate \n",
    "attention = output.permute(0, 2, 1, 3).reshape(batch_size, context_size, d_model)\n",
    "\n",
    "# Get output\n",
    "Wo = nn.Linear(d_model, d_model)\n",
    "output = Wo(attention)\n",
    "print(output.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "74f0b17b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply residual connection\n",
    "output = output + x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea5a7273",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply layer normalization\n",
    "layer_norm = nn.LayerNorm(d_model)\n",
    "layer_norm_output = layer_norm(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6df87b08",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply feed forward network\n",
    "output = nn.Linear(d_model, 4*d_model)(layer_norm_output)\n",
    "output = nn.ReLU()(output)\n",
    "output = nn.Linear(4*d_model, d_model)(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "593a9580",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 16, 64])\n"
     ]
    }
   ],
   "source": [
    "# Apply residual connection + LayerNorm\n",
    "output = output + layer_norm_output\n",
    "layer_norm = nn.LayerNorm(d_model)\n",
    "output = layer_norm(output)\n",
    "\n",
    "print(output.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a89a9bd3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 16, 199854])\n"
     ]
    }
   ],
   "source": [
    "# Define final linear layer\n",
    "output = nn.Linear(d_model, max_token_value + 1)(output)    # 每个词的维度展开 -> 每个词预测的下一个词的概率\n",
    "print(output.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "ca6c43e9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([6.0390e-06, 4.8566e-06, 6.1546e-06,  ..., 1.4018e-05, 2.5888e-06,\n",
      "        1.2664e-05], grad_fn=<SelectBackward0>)\n",
      "5.970880374661647e-05\n",
      "188862\n",
      " perioden\n"
     ]
    }
   ],
   "source": [
    "logits = F.softmax(output, dim=-1)\n",
    "predicted_index = torch.argmax(logits[0,0], dim=-1).item()\n",
    "\n",
    "print(logits[0, 0])\n",
    "print(max(logits[0, 0]).item())\n",
    "\n",
    "print(predicted_index)\n",
    "print(enc.decode([predicted_index]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "341d6763",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm-torch2.8-cu12.9-py3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
